{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aebdc562-1274-47fe-ab8e-4559b4ac9598",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn.functional as F\n",
    "\n",
    "device = torch.device(\"cuda:0\" if (torch.cuda.is_available()) else \"cpu\")\n",
    "\n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, data_dir, label_dir):\n",
    "        self.data_dir = data_dir\n",
    "        self.label_dir = label_dir\n",
    "        self.list_len = len(os.listdir(data_dir))\n",
    "        self.data_files = [f'{i}_data.npy' for i in range(self.list_len)]\n",
    "        self.label_files = [f'{i}_lab.npy' for i in range(self.list_len)]\n",
    "    def __len__(self):\n",
    "        return self.list_len\n",
    "    def __getitem__(self, idx):\n",
    "        data_file = self.data_files[idx]\n",
    "        label_file = self.label_files[idx]\n",
    "        data = np.load(os.path.join(self.data_dir, data_file))\n",
    "        label = np.load(os.path.join(self.label_dir, label_file))[0]\n",
    "        return torch.tensor(data, dtype=torch.float32).unsqueeze(0), torch.tensor(label, dtype=torch.long)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "58e3edf4-693c-4e07-b7a4-97e8fc932329",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Faildetect_CNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Faildetect_CNN, self).__init__()\n",
    "        self.conv1 = nn.Sequential(         \n",
    "            nn.Conv2d(in_channels=1,out_channels=16,kernel_size=3,stride=1,padding=1),                              \n",
    "            nn.ReLU(),                      \n",
    "            nn.MaxPool2d(kernel_size=2),    \n",
    "        )\n",
    "        self.conv2 = nn.Sequential(         \n",
    "            nn.Conv2d(16, 32, 3, 1, 1),     \n",
    "            nn.ReLU(),                      \n",
    "            nn.MaxPool2d(2),                \n",
    "        )\n",
    "        # fully connected layer, output 10 classes\n",
    "        self.FC1 = nn.Linear(32 * 45 * 4, 256)\n",
    "        self.FC2 = nn.Linear(256, 5)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "\n",
    "        x = x.view(x.size(0), -1) # flatten the output of conv2 to (batch_size, 32 * 7 * 7)      \n",
    "        x = F.relu(self.FC1(x))\n",
    "        output = self.FC2(x)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ccb42d02-a9f0-4908-a9f2-838c0cd1c75c",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data_dir = 'Train_data'\n",
    "train_label_dir = 'Train_lab'\n",
    "train_dataset = CustomDataset(train_data_dir, train_label_dir)\n",
    "train_loader = DataLoader(train_dataset, batch_size=5, shuffle=True)\n",
    "\n",
    "test_data_dir = 'Test_data'\n",
    "test_label_dir = 'Test_lab'\n",
    "test_dataset = CustomDataset(test_data_dir, test_label_dir)\n",
    "test_loader = DataLoader(test_dataset, batch_size=5, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f67738cf-a059-4bc6-a706-a8cd84a71f4c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Faildetect_CNN(\n",
      "  (conv1): Sequential(\n",
      "    (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU()\n",
      "    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (conv2): Sequential(\n",
      "    (0): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): ReLU()\n",
      "    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (FC1): Linear(in_features=5760, out_features=256, bias=True)\n",
      "  (FC2): Linear(in_features=256, out_features=5, bias=True)\n",
      ")\n",
      "# Model parameters: 1480901\n",
      "Epoch:0, train_loss:1.25161, train_acc:0.46564, test_loss:0.81825, test_acc:0.72245\n",
      "-----------------------------------------------\n",
      "Epoch:1, train_loss:0.66824, train_acc:0.74872, test_loss:0.73141, test_acc:0.70612\n",
      "-----------------------------------------------\n",
      "Epoch:2, train_loss:0.40141, train_acc:0.85744, test_loss:0.39096, test_acc:0.87755\n",
      "-----------------------------------------------\n",
      "Epoch:3, train_loss:0.26550, train_acc:0.90667, test_loss:0.29091, test_acc:0.91020\n",
      "-----------------------------------------------\n",
      "Epoch:4, train_loss:0.18426, train_acc:0.93436, test_loss:0.25732, test_acc:0.90204\n",
      "-----------------------------------------------\n",
      "Epoch:5, train_loss:0.12731, train_acc:0.95487, test_loss:0.30600, test_acc:0.88980\n",
      "-----------------------------------------------\n",
      "Epoch:6, train_loss:0.08208, train_acc:0.97744, test_loss:0.24581, test_acc:0.91020\n",
      "-----------------------------------------------\n",
      "Epoch:7, train_loss:0.06189, train_acc:0.97436, test_loss:0.24558, test_acc:0.90612\n",
      "-----------------------------------------------\n",
      "Epoch:8, train_loss:0.04363, train_acc:0.98769, test_loss:0.26138, test_acc:0.92245\n",
      "-----------------------------------------------\n",
      "Epoch:9, train_loss:0.01590, train_acc:0.99590, test_loss:0.21748, test_acc:0.93061\n",
      "-----------------------------------------------\n",
      "Finished Training\n"
     ]
    }
   ],
   "source": [
    "model = Faildetect_CNN()\n",
    "print(model)\n",
    "print('# Model parameters:', sum(param.numel() for param in model.parameters()))\n",
    "Criterion = nn.CrossEntropyLoss()\n",
    "Optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "train_loss = []\n",
    "train_acc = []\n",
    "test_loss = []\n",
    "test_acc  = []\n",
    "\n",
    "for epoch in range(10):\n",
    "    \n",
    "    train_loss_sum, test_loss_sum, train_num, test_num, train_i, test_i = 0.0, 0.0, 0, 0, 0, 0\n",
    "    train_acc_sum,test_acc_sum = 0, 0\n",
    "    TEST_acc_sum = 0\n",
    "    \n",
    "    for i_1, train_data in enumerate(train_loader):\n",
    "        \n",
    "        inputs, labels = train_data\n",
    "        inputs = inputs.type(torch.FloatTensor)\n",
    "        inputs = inputs.to(device)\n",
    "        \n",
    "        labels = labels.long()\n",
    "        labels = labels.to(device)\n",
    "        \n",
    "        model.train()\n",
    "        pre_labs = model(inputs)\n",
    "\n",
    "        Loss = Criterion(pre_labs, labels)\n",
    "\n",
    "        Optimizer.zero_grad()\n",
    "        Loss.backward()\n",
    "        Optimizer.step()\n",
    "        \n",
    "        train_loss_sum += Loss.item()\n",
    "        train_acc_sum += (pre_labs.argmax(dim=1) == labels).sum().item()\n",
    "        train_num += labels.shape[0]\n",
    "        train_i += 1\n",
    "\n",
    "    for i_2, data in enumerate(test_loader):\n",
    "        with torch.no_grad():\n",
    "\n",
    "            test_data, test_lab = data  \n",
    "            test_data = test_data.type(torch.FloatTensor)\n",
    "            test_data = test_data.to(device)\n",
    "            \n",
    "            test_lab = test_lab.long()\n",
    "            test_lab = test_lab.to(device)\n",
    "\n",
    "            model.eval()\n",
    "            pre_test = model(test_data)\n",
    "            t_loss = Criterion(pre_test, test_lab)\n",
    "\n",
    "            test_loss_sum += t_loss    \n",
    "            test_acc_sum += (pre_test.argmax(dim=1) == test_lab).sum().item()\n",
    "            test_num += test_lab.shape[0]                    \n",
    "            test_i += 1\n",
    "           \n",
    "    # if epoch >= 5:\n",
    "    #     torch.save(model, \"model/1/CNN_%d.pkl\"% epoch)\n",
    "    #     print(\"-------保存第%d epoch的模型---------\"% epoch)\n",
    "       \n",
    "    Train_Loss = train_loss_sum/train_i\n",
    "    Test_Loss = test_loss_sum/test_i\n",
    "    Train_ACC = train_acc_sum/train_num\n",
    "    Test_ACC = test_acc_sum/test_num\n",
    "\n",
    "    print('Epoch:%d, train_loss:%.5f, train_acc:%.5f, test_loss:%.5f, test_acc:%.5f' % \n",
    "          (epoch, Train_Loss, Train_ACC, Test_Loss, Test_ACC))\n",
    "    print('-----------------------------------------------')\n",
    "\n",
    "    train_loss.append(Train_Loss)            \n",
    "    train_acc.append(Train_ACC) \n",
    "    test_loss.append(Test_Loss) \n",
    "    test_acc.append(Test_ACC)       \n",
    "\n",
    "print(\"Finished Training\")\n",
    "\n",
    "# num_epochs = 10\n",
    "# for epoch in range(num_epochs):\n",
    "#     model.train()\n",
    "#     running_loss = 0.0\n",
    "#     for data, labels in train_loader:\n",
    "#         optimizer.zero_grad()\n",
    "#         outputs = model(data)\n",
    "#         loss = criterion(outputs, labels)\n",
    "#         loss.backward()\n",
    "#         optimizer.step()\n",
    "#         running_loss += loss.item()\n",
    "#     print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fad4a04c-0c8b-4873-8309-83041506e29a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
